"""
write by qianqianjun
2020.02.15
准备数据集
"""
import os
import numpy as np

class PrepareData:
    def __init__(self,dir,shuffle=True,include_subdirectory=True,random_seed=19970329):
        """
        数据集准备工具类
        从数据集目录中提取数据的地址保存到一个列表中
        :param dir:  数据集的目录
        :param shuffle:  是否打乱顺序，默认打乱
        :param include_subdirectory: 在提取数据的时候，是否把dir 目录中子目录中的图片也包含进去
        :param random_seed: 设置随机种子
        """
        self.rootDir=dir
        self.seed=random_seed
        images_address=[]

        if include_subdirectory:
            for current_dir,subdir,files in os.walk(self.rootDir):
                images_addr_set=self.removeOtherFile(current_dir,files)
                images_address.extend(images_addr_set)
        else:
            files=os.listdir(self.rootDir)
            images_addr_set=self.removeOtherFile(self.rootDir,files)
            images_address.extend(images_addr_set)
        self.images_address=np.array(images_address)

        self.length = len(self.images_address)
        if shuffle:
            self.shuffle(seed=random_seed)


    def removeOtherFile(self,curdir:str,files:list)->list:
        """
        移除一个列表中不是图片文件的其它文件
        :param curdir: 当前地址
        :param files:  文件地址列表
        :return: 图片文件的地址列表
        """
        image_arr=[]
        for item in files:
            if item.endswith("jpg"):
                image_arr.append(os.path.join(curdir,item))
        return image_arr

    def shuffle(self,seed):
        """
        用于对数据进行重新排序
        :param seed: 随机种子，为了确保实验的可重复性
        :return:
        """
        np.random.seed(seed)
        index=np.random.permutation(self.length)
        self.images_address=self.images_address[index]

    def train_test_split(self,total=None,train:int=None):
        """
        分割训练和测试数据集
        :param total:  数据集总数
        :param train:  训练数据所占比重
        :return:  训练集和测试集图片的地址列表
        """
        if total==None:
            total=self.length
        if train==None:
            train=self.length/2
        assert total<=self.length ,"共有图片{0}张,不够{1}张！".format(self.length,total)
        train_num=train
        test_num=total-train_num

        train_set=self.images_address[:train_num]
        test_set=self.images_address[train_num:train_num+test_num]

        return train_set,test_set

    def save_list(self,save_dir:str,train_set:np.array,test_set:np.array):
        """
        将分割好的数据集进行保存
        :param save_dir:  保存数据集的地址
        :param train_set:  要保存的训练集
        :param test_set:  要保存的测试集
        :return:
        """
        if not os.path.exists(save_dir):
            os.makedirs(save_dir,exist_ok=True)
        import pickle as pk
        with open(os.path.join(save_dir,"trainSet"),"wb") as f:
            pk.dump(train_set,f)
        with open(os.path.join(save_dir,"testSet"),'wb') as f:
            pk.dump(test_set,f)

        """
        测试代码
        print(train_set)
        print("----------------")
        print(test_set)
        print("----------------")
        with open(os.path.join(save_dir,"trainSet"),mode="rb") as f:
            train_set=pk.load(f)
        with open(os.path.join(save_dir,"testSet"),mode="rb") as f:
            test_set=pk.load(f)

        print(train_set)
        print("----------------")
        print(test_set)
        """

import torch.utils.data as data
from PIL import Image
class DataSet(data.Dataset):
    def __init__(self,files,resizeTo:int=64):
        """
        创建数据集
        :param files: 图片文件地址列表
        :param resizeTo:  将图片转换为多大
        """
        self.files=files
        self.resizeTo=resizeTo
        self.length=len(files)

    def __getitem__(self, index):
        """
        重载父类的方法
        :param index:
        :return:
        """
        filepath=self.files[index]
        image=self.imageFileReader(filepath,resizeTo=self.resizeTo)
        return  image

    def __len__(self):
        return self.length

    def imageFileReader(self,filepath:str,resizeTo:int):
        """
        图片文件读取方法
        :param filepath: 图片文件地址
        :param resizeTo: 像图片转换到多大（边长）
        :return: 返回图片的 np.array() 数组表示
        """
        img=Image.open(filepath)
        img=img.resize((resizeTo,resizeTo),Image.ANTIALIAS)
        return np.array(img)